#include "common.h"
#include "paging.h"

#define PANIC(msg) panic(msg, __FILE__, __LINE__);
extern uint end;
uint placement_address = (uint)&end;

page_directory_t *kernel_directory = 0;
page_directory_t *current_directory = 0;

uint kmalloc(uint size){
	uint tmp = placement_address;
	placement_address += size;
}

uint kmalloc_a(uint size){
		// align it.
		placement_address &= 0xFFFFF000;
		placement_address += 0x1000;
	uint tmp = placement_address;
	placement_address += size;
	return tmp;
}	

uint kmalloc_p(uint size, uint *phys){
	uint temp = placement_address;
	*phys = placement_address;
	placement_address += size;
	return temp;	
}

uint kmalloc_ap(uint size, uint *phys){
	placement_address &= 0xFFFFF000;
	placement_address += 0x1000;
	if(phys){
		*phys = placement_address;
	}
	uint temp = placement_address;
	placement_address += size;
	return temp;
}

// A bitset of frames - used or free.
uint *frames;
uint nframes;

// Defined in kheap.c
extern uint placement_address;

// Macros used to access the bitset algorithms.
#define INDEX_FROM_BIT(a) (a/(8*4))
#define OFFSET_FROM_BIT(a) (a%(8*4))

// Static function to set a bit in the frames bitset.
static void set_frame(uint frame_addr){
	uint frame = frame_addr / 0x1000;
	uint idx = INDEX_FROM_BIT(frame);
	uint off = OFFSET_FROM_BIT(frame);
	frames[idx] |= (0x1 << off);
}

// Static function to clear a bit int the frames bitset
static void clear_frame(uint frame_addr){
	uint frame = frame_addr/0x1000;
	uint idx = INDEX_FROM_BIT(frame);
	uint off = OFFSET_FROM_BIT(frame);
	frames[idx] &= ~(0x1 << off);
}

// Static function to find the first free frame
static uint test_frame(uint frame_addr){
	uint frame = frame_addr/0x1000;
	uint idx = INDEX_FROM_BIT(frame);
	uint off = OFFSET_FROM_BIT(frame);
	return (frames[idx] & (0x1 << off));
}

// Static function to find the first free frame.
static uint first_free_frame(){
	uint i, j;
	for(i = 0; i < INDEX_FROM_BIT(nframes); i++){
		if(frames[i] != 0xFFFFFFFF){ // Nothing Free, exit early
			// at least One bit is free.
			for(j = 0; j < 32; j++){
				uint toTest = 0x1 << j;
				if( !(frames[i]&toTest) ){
					return i*8*4 + j;
				}
			}
		}
	}
}

// Function to allocate Frame.
void alloc_frame(page_t *page, int is_kernel, int is_writable){
	if (page->frame!= 0)
		return;	// Frame was already allocated. return straight away.
	else{
		uint idx = first_free_frame();
		if (idx == (uint)-1){
			PANIC("No Free Frames !");
		}
		set_frame(idx*0x1000);	// This frame is now ours.
		page->present = 1;
		page->rw = (is_writable)?1:0;
		page->user = (is_kernel)?0:1;
		page->frame = idx;
	}
}

// Function to deallocate a frame
void free_frame(page_t *page){
	uint frame;
	if (!(frame=page->frame)){
		return;		// The page didn't actually have an allocated frame.
	}else{
		clear_frame(frame);	// Frame is now free again.
		page->frame=0x0;	// Page now doesnt have a frame.
	}
}

// Paging Code.
void init_paging(){
	// The size of the physical memory. For the moment, we assume
	// it is 16mb.
	uint mem_end_page = 0x1000000;
	
	nframes = mem_end_page / 0x1000;
	frames = (uint*)kmalloc(INDEX_FROM_BIT(nframes));
	memset(frames, 0, INDEX_FROM_BIT(nframes));
	
	// Let's Make a page directory.
	kernel_directory = (page_directory_t *)kmalloc_a(sizeof(page_directory_t));
	memset(kernel_directory, 0, sizeof(page_directory_t));
	current_directory = kernel_directory;
	
	
	/* We need to identify map (phys adr = virt addr) from 0x0 to the 
	 * the end of used memory, so we can access this transparently,
	 * as if paging wasn't enabled.
	 * NOTE that we use a while loop here deliberately.
	 * Inside the loop we actually change placement_address
	 * by calling kmalloc(). A while loop causes this to be computed
	 * on the fly rather than once at the start.
	 */
	int i = 0;
	while(i < placement_address){
		// Kernel code is readable but not writable from userspace.
		alloc_frame(get_page(i, 1, kernel_directory), 0, 0);
		i += 0x1000;
	}
	// Before we enable paging, we must register our page fault handler.
	register_interrupt_handler(14, page_fault);
	
	// Now enable paging!
	switch_page_directory(kernel_directory);	
}

void switch_page_directory(page_directory_t *dir){
	current_directory = dir;
	asm volatile("mov %0, %%cr3"::"r"(&dir->tablesPhysical));
	uint cr0;
	asm volatile("mov %%cr0, %0":"=r"(cr0));
	cr0 |= 0x80000000;	// Enable Paging.
	asm volatile("mov %0, %%cr0"::"r"(cr0));
}

page_t *get_page(uint address, int make, page_directory_t *dir){
	// Turn the address ino the index.
	address /= 0x1000;
	// Find the page table containing this address.
	uint table_idx = address / 1024;
	if(dir->tables[table_idx]) // If this page is already assigned.
		return &dir->tables[table_idx]->pages[address%1024];
	else if(make){
		uint temp;
		dir->tables[table_idx] = (page_table_t *)kmalloc_ap(sizeof(page_table_t), &temp);
		memset(dir->tables[table_idx], 0, 0x1000);
		dir->tablesPhysical[table_idx] = temp | 0x7; // Presend RW, USER.
		return &dir->tables[table_idx]->pages[address%1024];
	}
	else
	return 0;
}

void page_fault(registers_t regs){
	// A page fault has occured.
	// The faulting address is stored in the cr2 register
	uint faulting_address;
	asm volatile("mov %%cr2, %0" : "=r" (faulting_address));
	
	// The error code gives us details of what happened.
	int present = !(regs.err_code &0x1); // Page not present.
	int rw = regs.err_code & 0x2;		// Write Operation ?
	int us = regs.err_code & 0x4;		// Processor was in user mode ?
	int reserved = regs.err_code & 0x8;	// Overwritten CPU-reserved bits of a page.
	int id = regs.err_code & 0x10;		// Caused by an instruction fetch.
	
	// Output an error message.
	printk("PageFault! (");
	if(present) printk("present ");
	if(rw) printk("read-only ");
	if(us) printk("user-mode ");
	if(reserved) printk("reserved ");
	printk(") at 0x%x\n", faulting_address);
	PANIC("Page Fault!");
	
}

